import pickle
from common.mnist import load_mnist
# PIL python image library 是一个模块
from PIL import Image
import forward
import numpy as np


# 显示图片
def show_image(img):
    pil_img = Image.fromarray(np.uint8(img))
    pil_img.show()


#   (训练图像, 训练标签), (测试图像, 测试标签)
# 返回了测试图像和测试标签，numpy形式
def get_data():
    (x_train, t_train), (x_test, t_test) = load_mnist(normalize=True, flatten=True, one_hot_label=False)
    return x_test, t_test


def init_network():
    # 书上没有说这个pkl是怎么来的，在推理阶段，我们直接使用这些参数
    with open("sample_weight.pkl", "rb") as f:
        network = pickle.load(f)
    return network


if __name__ == '__main__':
    # 批处理可以加速，因为库有优化。
    batch_size = 100
    x, t = get_data()
    network = init_network()
    accuracy_cnt = 0
    # 这里的for循环i的范围是从0到x.size ，每次增加 batch_size
    for i in range(0, len(x), batch_size):
        x_batch = x[i:i + batch_size]
        y_batch = forward.forword(network, x_batch)
        # 这里的y_batch 的维度为 100*10
        p = np.argmax(y_batch, axis=1)  # 沿着第一维的方向
        # 这里p的维度为 100*1，十个数中只保留最大的那一个。
        if i == 0:
            print(y_batch.shape)
            print(x.shape)
            print(x_batch.shape)
            print(p.shape)
        accuracy_cnt += np.sum(p == t[i:i + batch_size])
    print("Accuracy:" + str(float(accuracy_cnt) / len(x)))
